package top.lywivan.tree;

import top.lywivan.printer.BinaryTreeInfo;

/**
 * @desc
 * @author:ivan.liu
 */
public class RBTree implements BinaryTreeInfo {

    private static final boolean RED=false;
    private static final boolean BLACK=true;

    //定义根节点
    Node root;

    /**
     * 给红黑树添加节点
     * @param key
     */
    public void add(int key){
        //先正常为BST添加节点
        if(root==null){
            root=new Node(key,null);
            afterAdd(root);
            return;
        }
        Node parent=null;
        Node curr=root;
        while (curr!=null){
            if(key>curr.key){
                parent=curr;
                curr=curr.right;
            }else if(key<curr.key){
                parent=curr;
                curr=curr.left;
            }else {
                //相等时不做处理
            }
        }
        //为parent添加一个子节点
        Node newNode=new Node(key,parent);
        if(key> parent.key){
            parent.right=newNode;
        }else if(key< parent.key) {
            parent.left=newNode;
        }

        //已经添加完成，进行红黑树的自平衡
        afterAdd(newNode);
    }

    /**
     * 添加节点后处理红黑树自平衡
     * @param node 新添加的节点
     */
    public void afterAdd(Node node){
        //获取父节点
        Node parent = node.parent;
        //父节点为空，则将当前节点变为黑色
        if(parent==null){
            black(node);
            return;
        }
        //父节点为黑色，不需要进行额外处理
        if(isBlack(parent)){
            return;
        }
        //父节点为红色，则根据叔父节点的颜色处理
        Node uncle=parent.sibling();
        Node grand=parent.parent;
        //叔父节点是红色，则只染色，不旋转。需要进行递归处理
        if(isRed(uncle)){
            //父节点，叔父染黑色，祖父染红
            black(parent);
            black(uncle);
            red(grand);
            //祖父节点当作新添加的节点，继续处理
            afterAdd(grand);
            return;
        }
        //叔父节点是黑色，则需要区分RR，LL，RL，LR
        if(parent.isLeftChild()){
            if(node.isLeftChild()){
                //LL的情况，父节点染黑色，祖父染红色
                black(parent);
                red(grand);
                //祖父节点LL右旋转
                LLRotate(grand);
            }else {
                //LR的情况，当前节点染黑，祖父节点染红
                black(node);
                red(grand);
                //祖父节点LR旋转
                LRRotate(grand);
            }
        }else {
            if(node.isLeftChild()){
                //RL的情况，当前节点染黑，祖父节点染红
                black(node);
                red(grand);
                //祖父节点RL旋转
                RLRotate(grand);
            }else {
                //RR的情况，父节点染黑色，祖父染红色
                black(parent);
                red(grand);
                //祖父节点RR左旋转
                RRRotate(grand);
            }
        }
    }

    /**
     * 删除节点
     * @param key
     */
    public void remove(int key){
        //找到需要删除的节点
        Node node=node(key);
        //删除具体的节点
        remove(node);
    }

    /**
     * 删除具体的节点
     * @param node
     */
    private void remove(Node node) {
        if (node==null){
            return;
        }
        //处理度为2的节点，采用前驱代替
        if (node.hasTwoChildren()){
            //找到前驱节点，即左子树的最大节点
            Node p=node.left;
            while (p.right!=null){
                p=p.right;
            }
            //替换前驱节点
            node.key=p.key;
            node=p;
        }
        //删除度为1或者度为0的节点
        //处理度为1的节点
        Node replace=node.left!=null?node.left:node.right;
        //判断删除节点是否是根节点
        if(node.parent==null){
            root=replace;
        }else if (node.isLeftChild()){
            node.parent.left=replace;
        }else if (node.isRightChild()){
            node.parent.right=replace;
        }
        //处理parent指针
        if(replace!=null){
            replace.parent=node.parent;
        }
        //调整红黑树
        afterRemove(node,replace);
        //清空node的指针
        node.parent=null;
        node.left=null;
        node.right=null;
    }

    /**
     * 删除时，调整红黑树的平衡
     * @param beRemoved
     * @param replace
     */
    private void afterRemove(Node beRemoved, Node replace) {
        //被删除的是红色，则无需调整
        if(isRed(beRemoved)){
            return;
        }
        //被删除的是黑色节点
        //如果有一个红色的替代子节点，则将其染黑即可
        if(isRed(replace)){
            black(replace);
            return;
        }
        //此时被删除节点为黑色，并且度为0
        Node parent = beRemoved.parent;
        //如果是根节点，则不需要处理
        if(parent==null){
            return;
        }
        //判断被删除节点是左边还是右边。需要考虑递归调用的情况
        boolean left=parent.left==null||beRemoved.isLeftChild();
        //获取兄弟节点
        Node sibling=left?parent.right:parent.left;
        //分别处理左和右的情况
        if(left){
            //处理兄弟是红色的情况，需要换兄弟
            if(isRed(sibling)){
                //兄弟染黑色，父节点染红色，父节点RR选择
                black(sibling);
                red(parent);
                RRRotate(parent);
                //更换旋转后的最新的兄弟节点
                sibling=parent.right;
            }
            //处理兄弟是黑色的情况，需要看兄弟是否有红色子节点
            if (isBlack(sibling.left)&&isBlack(sibling.right)){
                //如果父节点本身就是黑色，则需要递归处理
                boolean parentBlack=isBlack(parent);
                //此时只染色，兄弟染红色，父节点染黑色
                red(sibling);
                black(parent);
                if(parentBlack){
                    afterRemove(parent,null);
                }
            }else {
                //此时需要通过旋转借一个红色子节点
                Node newParent;
                //兄弟的右孩子是红色，则RR旋转
                if (isRed(sibling.right)){
                    newParent = RRRotate(parent);
                }else{
                    //兄弟的右孩子是黑色，则RL旋转
                    newParent=RLRotate(parent);
                }
                //父节点继承原来父节点颜色
                color(newParent,colorOf(parent));
                //左右子节点染黑色
                black(newParent.left);
                black(newParent.right);
            }
        }else {
            //处理兄弟是红色的情况，需要换兄弟
            if(isRed(sibling)){
                //兄弟染黑色，父节点染红色，父节点RR选择
                black(sibling);
                red(parent);
                LLRotate(parent);
                //更换旋转后的最新的兄弟节点
                sibling=parent.left;
            }
            //处理兄弟是黑色的情况，需要看兄弟是否有红色子节点
            if (isBlack(sibling.left)&&isBlack(sibling.right)){
                //如果父节点本身就是黑色，则需要递归处理
                boolean parentBlack=isBlack(parent);
                //此时只染色，兄弟染红色，父节点染黑色
                red(sibling);
                black(parent);
                if(parentBlack){
                    afterRemove(parent,null);
                }
            }else {
                //此时需要通过旋转借一个红色子节点
                Node newParent;
                //兄弟的右孩子是红色，则RR旋转
                if (isRed(sibling.left)){
                    newParent = LLRotate(parent);
                }else{
                    //兄弟的右孩子是黑色，则RL旋转
                    newParent=LRRotate(parent);
                }
                //父节点继承原来父节点颜色
                color(newParent,colorOf(parent));
                //左右子节点染黑色
                black(newParent.left);
                black(newParent.right);
            }
        }
    }

    /**
     * 根据key找到要删除的节点
     * @param key
     * @return
     */
    private Node node(int key) {
        Node curr=root;
        while (curr!=null){
            if(key>curr.key){
                curr=curr.right;
            }else if (key<curr.key){
                curr=curr.left;
            }else {
                return curr;
            }
        }
        return null;
    }

    /**
     * RR左旋转，返回旋转后的根节点
     * @param node
     * @return
     */
    public Node RRRotate(Node node){
        //先获取新的根节点
        Node newRoot=node.right;
        //处理原来的左子节点
        node.right=newRoot.left;
        if(newRoot.left!=null){
            newRoot.left.parent=node;
        }
        //处理原来的根节点，作为新根节点的子节点
        newRoot.left=node;
        newRoot.parent=node.parent;
        //RR是在右子树的右边，但node不一定在其父节点的右边，因此需要进行判断
        if (node.isLeftChild()){
            node.parent.left=newRoot;
        }else if (node.isRightChild()){
            node.parent.right=newRoot;
        }else {
            //node为根节点，则更新新的根节点
            root=newRoot;
        }
        //修改node的父亲
        node.parent=newRoot;
        return newRoot;
    }

    /**
     * LL右旋转，返回旋转后的根节点
     * @param node
     * @return
     */
    public Node LLRotate(Node node){
        //先获取新的根节点
        Node newRoot=node.left;
        //处理原来的右子节点
        node.left=newRoot.right;
        if(newRoot.right!=null){
            newRoot.right.parent=node;
        }
        //处理原来的根节点，作为新根节点的子节点
        newRoot.right=node;
        newRoot.parent=node.parent;
        //RR是在右子树的右边，但node不一定在其父节点的右边，因此需要进行判断
        if (node.isLeftChild()){
            node.parent.left=newRoot;
        }else if (node.isRightChild()){
            node.parent.right=newRoot;
        }else {
            //node为根节点，则更新新的根节点
            root=newRoot;
        }
        //修改node的父亲
        node.parent=newRoot;
        return newRoot;
    }

    /**
     * RL旋转，右子节点先右旋，当前节点再左旋
     * @param node
     * @return
     */
    public Node RLRotate(Node node){
        LLRotate(node.right);
        return RRRotate(node);
    }

    /**
     * LR旋转，左子节点先左旋，当前节点再右旋
     * @param node
     * @return
     */
    public Node LRRotate(Node node){
        RRRotate(node.left);
        return LLRotate(node);
    }

    /**
     * 给节点染色，并返回该节点
     * @param node
     * @param color
     * @return
     */
    public Node color(Node node,boolean color){
        if(node==null){
            return null;
        }
        node.color=color;
        return node;
    }

    /**
     * 直接将节点染黑色
     * @param node
     * @return
     */
    public Node black(Node node){
        return color(node,BLACK);
    }

    /**
     * 直接将节点染红色
     * @param node
     * @return
     */
    public Node red(Node node){
        return color(node,RED);
    }

    /**
     * 获取当前节点的颜色
     * 空节点应该是黑色
     * @param node
     * @return
     */
    public boolean colorOf(Node node){
        return node==null?BLACK:node.color;
    }

    /**
     * 判断节点是否是黑色
     * @param node
     * @return
     */
    public boolean isBlack(Node node){
        return colorOf(node)==BLACK;
    }

    /**
     * 判断节点是否是红色
     * @param node
     * @return
     */
    public boolean isRed(Node node){
        return colorOf(node)==RED;
    }

    @Override
    public Object root() {
        return root;
    }

    @Override
    public Object left(Object node) {
        return ((Node)node).left;
    }

    @Override
    public Object right(Object node) {
        return ((Node)node).right;
    }

    @Override
    public Object string(Object node) {
        return node;
    }

    public static class Node {
        int key;
        //false为红色，true为黑色。新添加的节点都是红色
        boolean color=RED;
        Node left;
        Node right;
        //需要存储当前节点的父节点，便于进行处理
        Node parent;

        public Node(int key,Node parent){
            this.key=key;
            this.parent=parent;
        }

        /**
         * 判断当前节点是否有两个子节点
         * @return
         */
        public boolean hasTwoChildren(){
            return left!=null&&right!=null;
        }

        /**
         * 判断当前节点是否是其父节点的左子孩子
         * @return
         */
        public boolean isLeftChild(){
            return parent!=null&&parent.left==this;
        }

        /**
         * 判断当前节点是否是其父节点的右子孩子
         * @return
         */
        public boolean isRightChild(){
            return parent!=null&&parent.right==this;
        }

        /**
         * 返回当前节点的兄弟节点
         * @return
         */
        public Node sibling(){
            if(isLeftChild()){
                return parent.right;
            }
            if(isRightChild()){
                return parent.left;
            }
            return null;
        }

        @Override
        public String toString() {
            String str="";
            if(color==RED){
                str="R_";
            }
            return str+key;
        }
    }
}
